package com.lili.collection.mq.plugins;

import com.lili.collection.core.context.BatchRecordContext;
import com.lili.collection.core.context.StageContext;
import com.lili.collection.core.domain.RunnerMsgTransfer;
import com.lili.collection.core.enums.DataFormat;
import com.lili.collection.core.error.OutputRunnerException;
import com.lili.collection.core.error.StageRecordException;
import com.lili.collection.core.plugins.OutputPlugin;
import com.lili.collection.core.record.Record;
import com.lili.collection.core.runner.result.ErrOutputData;
import com.lili.collection.core.runner.result.RecordInfo;
import com.lili.collection.core.runner.result.RunnerResult;
import com.lili.collection.core.utils.HandlerUtil;
import com.lili.collection.core.utils.NoticeUtil;
import com.lili.collection.mq.config.RabbitmqOutputConfig;
import com.lili.collection.mq.domain.ByteMessageWrapper;
import com.lili.collection.mq.utils.JsonConvertUtil;
import com.lili.collection.mq.utils.MQUtil;
import com.rabbitmq.client.*;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 描述: rabbitmq消息队列输出插件
 *
 * @date:2023/5/23 16:00
 * @author:lili
 */
@Slf4j
public class RabbitmqOutputPlugin implements OutputPlugin<RabbitmqOutputConfig> {
    //最大消息限制为4M
    private static final int MAX_MESSAGE_SIZE = 4 * 1024 * 1024;
    private static Map<String, Return> returnMap = new ConcurrentHashMap<>();

    private RabbitmqOutputConfig config;

    private Throwable tempThrowable = null;
    private boolean sourceHasValue = false;
    private int executeNum;
    @Getter
    private RunnerResult runnerResult;
    @Getter
    private RunnerMsgTransfer transfer;
    private Connection connection;
    private Channel channel;
    private String routingKey = "";
    private String exchangeName;
    private String instanceId;
    private boolean connClosed;

    @Override
    public void setRunnerResult(RunnerResult outputRunnerResult) {
        this.runnerResult = outputRunnerResult;
    }

    @Override
    public void setRunnerTransfer(RunnerMsgTransfer transfer) {
        this.transfer = transfer;
    }

    @Override
    public String queryInitialOffset() {
        return this.config.getInitialOffset();
    }

    @Override
    public void preHandle() {
        //进行维护队列与交换机之间的绑定关系
        if (!this.config.isQueueExchangeBinded()) {
            log.info("不维护队列与交换机间的绑定!");
            return;
        }
        log.info("开启维护队列与交换机间的绑定!");
        String exchangeName = this.config.getExchangeName();
        List<String> queueList = this.config.getQueueList();
        if (CollectionUtils.isEmpty(queueList)) {
            log.error("队列为空,取消绑定!");
        }
        try {
            for (String queueName : queueList) {
                this.channel.queueBind(queueName, exchangeName, routingKey);
            }
        } catch (Exception e) {
            log.error("队列与交换机绑定错误!", e);
            throw new OutputRunnerException("队列" + queueList + "与交换机[" + exchangeName + "]绑定错误!");
        }
    }

    @Override
    public synchronized void execute(StageContext stageContext) {
        BatchRecordContext recordContext = (BatchRecordContext) stageContext;
        RecordInfo recordInfo = new RecordInfo();
        this.executeNum++;
        boolean isCommit = false;
        RunnerResult<String> outputRunnerResult = this.getRunnerResult();
        boolean isRollback = false;
        List<Record> output = recordContext.getOutput();
        try {
            try {
                if (this.getTransfer() != null && this.getTransfer().isRunError()) {
                    throw new IllegalStateException("Rabbitmq异步运行错误!");
                }
                if (output.size() == 0) {
                    if (returnMap.remove(instanceId) != null) {
                        throw new RuntimeException("消息投放队列失败!");
                    }
                    if (sourceHasValue && config.isEnableTransaction()) {
                        this.channel.txCommit();
                    }
                    if (!connClosed) {
                        connClosed = MQUtil.close(this.channel, this.connection);
                    }
                    log.info("-->所有数据写入完成,提交事务成功!");
                    return;
                }

                if (config.getDataFormat() == DataFormat.ORIGIN ||  config.getDataFormat()==DataFormat.ORIGIN_CONVERT) {
                    for (Record record : output) {
                        sendDataToMQ(String.valueOf(record.getField().getValue()));
                        recordInfo.incrementSuccessCount();
                    }
                } else {
                    //转换消息
                    List<ByteMessageWrapper> byteMessageWrappers = JsonConvertUtil.convertRecordToJsonByteWrapper(output, recordInfo, MAX_MESSAGE_SIZE, true);
                    for (ByteMessageWrapper byteMessageWrapper : byteMessageWrappers) {
                        if (executeNum == 1 && recordInfo.getSuccessRecordCount() > 0) {
                            sourceHasValue = true;
                        }
                        HandlerUtil.sendProcessedData(recordInfo,config, byteMessageWrapper.getOrginJsonData(), (data) -> {
                            this.sendDataToMQ(data);
                            return true;
                        });
                        //消息推送
                        channel.basicPublish(this.config.getExchangeName(), routingKey, true, false, MessageProperties.PERSISTENT_TEXT_PLAIN, byteMessageWrapper.getBytes());
                    }
                }
                if (config.isIncrementalMode()) {
                    if (returnMap.remove(instanceId) != null) {
                        throw new RuntimeException("消息投放队列失败!");
                    }
                    if (config.isEnableTransaction()) {
                        // 只有增量需要每批次提交事务
                        this.channel.txCommit();
                    }
                    isCommit = true;
                } else if (recordContext.getStageContextState() == StageContext.StageContextState.COMPLETE) {
                    if (returnMap.remove(instanceId) != null) {
                        throw new RuntimeException("消息投放队列失败!");
                    }
                    if (config.isEnableTransaction()) {
                        this.channel.txCommit();
                    }
                }
                log.info("-->第{}批次数据写入完成\r\n-->写入量：{}", recordContext.getIndex(), output.size());
            } catch (Throwable throwable) {
                if (config.isEnableTransaction()) {
                    this.channel.txRollback();
                }
                isRollback = true;
                throw throwable;
            } finally {
                if ((isRollback || recordContext.getStageContextState() == StageContext.StageContextState.COMPLETE) && !connClosed) {
                    connClosed = MQUtil.close(this.channel, this.connection);
                }
            }
        } catch (Throwable throwable) {
            tempThrowable = throwable;
            RunnerMsgTransfer transfer = this.getTransfer();
            if (transfer != null) {
                transfer.setRunError(true);
                transfer.interruptRunner();
            }
            if (config.isEnableTransaction()) {
                recordInfo.setErrorRecordCount(recordInfo.getSuccessRecordCount());
                recordInfo.setSuccessRecordCount(0);
            }
            recordInfo.setSuccessed(false);
            recordContext.setStageRecordState(StageContext.StageRecordState.ALL_ERROR);
            recordContext.setException(new StageRecordException(throwable));
            outputRunnerResult.setSuccessed(false);
            //错误日志收集
            ErrOutputData errOutputData = recordInfo.getErrOutputData();
            errOutputData.setBatchIndex(recordContext.getIndex());
            errOutputData.setErrorMessage(throwable.toString());
            errOutputData.setFirstBatchError(executeNum == 1 ? true : false);
            errOutputData.setTransaction(config.isEnableTransaction());
            errOutputData.setType("BATCH-OUTPUT-RABBITMQ");
            log.info("-->第{}批次数据写入时发生异常\r\n-->原因：{}",
                    recordContext.getIndex(),
                    throwable.getLocalizedMessage(), throwable);
        }
        if (tempThrowable != null) {
            ErrOutputData errOutputData = recordInfo.getErrOutputData();
            outputRunnerResult.setErrOutputData(errOutputData);
            if (!config.isIncrementalMode() && config.isEnableTransaction()) {
                this.clearResultCount(outputRunnerResult);
            } else if (!config.isEnableTransaction()) {
                this.addResultCount(recordInfo, outputRunnerResult);
            }
        } else {
            if (!config.isIncrementalMode()) {
                this.addResultCount(recordInfo, outputRunnerResult);
            }
            if (config.isIncrementalMode() && isCommit) {
                this.addResultCount(recordInfo, outputRunnerResult);
            }
        }
        recordContext.clearOutput();
    }

    private void sendDataToMQ(String data) throws Exception {
        if(StringUtils.isBlank(data)){
            throw new IllegalArgumentException("发送MQ数据不能为空！");
        }
        NoticeUtil.sendDataAndSaveNotice(config,data,(sendData)->{
            channel.basicPublish(this.config.getExchangeName(), routingKey, true, false, MessageProperties.PERSISTENT_TEXT_PLAIN, sendData.getBytes(StandardCharsets.UTF_8));
            return true;
        });
    }

    private void clearResultCount(RunnerResult<String> outputRunnerResult) {
        outputRunnerResult.setSuccessRecordCount(0);
        outputRunnerResult.setTotalRecordCount(0);
        outputRunnerResult.setErrorRecordCount(0);
    }

    private void addResultCount(RecordInfo recordInfo, RunnerResult<String> outputRunnerResult) {
        outputRunnerResult.addSuccessCount(recordInfo.getSuccessRecordCount());
        outputRunnerResult.addTotalCount(recordInfo.getTotalRecordCount());
        outputRunnerResult.addErrorCount(recordInfo.getErrorRecordCount());
    }

    @Override
    public void dispose() {
        if (!connClosed) {
            connClosed = MQUtil.close(this.channel, this.connection);
        }
    }

    @Override
    public void init(RabbitmqOutputConfig config) {
        if(this.config.getDataFormat()==null){
            throw new IllegalArgumentException("rabbitmq输入输出数据格式定义不能为空！");
        }
        this.instanceId = UUID.randomUUID().toString();
        this.config = config;
        this.exchangeName = this.config.getExchangeName();
        this.connection = this.createConnection();
        this.channel = this.createChannel();
        this.addReturnListener();
    }

    private void addReturnListener() {
        this.channel.addReturnListener(new ReturnCallback() {
            @Override
            public void handle(Return returnMessage) {
                String exchange = returnMessage.getExchange();
                if (StringUtils.isNotBlank(exchange) && exchange.equals(exchangeName)) {
                    returnMap.put(instanceId, returnMessage);
                }
            }
        });
    }

    private Channel createChannel() {
        try {
            Channel channel = this.connection.createChannel();
            //启用事务消息
            if (config.isEnableTransaction()) {
                channel.txSelect();
            }
            return channel;
        } catch (Exception e) {
            log.error("创建通道错误!", e);
            throw new OutputRunnerException(e);
        }
    }

    private Connection createConnection() {
        try {
            ConnectionFactory connectionFactory = new ConnectionFactory();
            connectionFactory.setVirtualHost(this.config.getVirtualHost());
            connectionFactory.setUsername(this.config.getUsername());
            connectionFactory.setPassword(this.config.getPassword());
            return connectionFactory.newConnection(this.config.getAddressList());
        } catch (Exception e) {
            log.error("创建连接错误!", e);
            throw new OutputRunnerException(e);
        }
    }
}